#include "test_controller/ssh.h"
#include <fstream>
#include "base/string_util.h"

SSH::SSH()
  : session_(NULL) {
}


SSH::~SSH() {
  Uninitialize();
  std::cout << "Realse SSH" << std::endl;
}

bool SSH::Initialize(
  const char *host,
  int port,
  const char *username,
  const char *password) {
  if (host == NULL) {
    fprintf(stderr, "host can not be null");
    return false;
  }

  if (username == NULL) {
    fprintf(stderr, "username can not be null");
    return false;
  }

  if (password == NULL) {
    fprintf(stderr, "password can not be null");
    return false;
  }

  session_ = ssh_new();
  if (NULL == session_) {
    return false;
  }

  if (ssh_options_set(session_, SSH_OPTIONS_HOST, host) < 0) {
    ssh_free(session_);
    session_ = NULL;
    return false;
  }

  if (ssh_options_set(session_, SSH_OPTIONS_USER, username) < 0) {
    ssh_free(session_);
    session_ = NULL;
    return false;
  }

  if (ssh_options_set(session_, SSH_OPTIONS_PORT, &port) < 0) {
    ssh_free(session_);
    session_ = NULL;
    return false;
  }

  if (ssh_connect(session_)) {
    fprintf(stderr, "Connection failed : %s\n", ssh_get_error(session_));
    ssh_disconnect(session_);
    ssh_free(session_);
    session_ = NULL;
    return false;
  }

  int rc = ssh_userauth_password(session_, username, password);
  if (rc != SSH_AUTH_SUCCESS) {
    fprintf(stderr, "Error authenticating with password: %s\n",
            ssh_get_error(session_));
    ssh_disconnect(session_);
    ssh_free(session_);
    session_ = NULL;
    return false;
  }

  shell_channel_ = ssh_channel_new(session_);;
  if (shell_channel_ == NULL) {
    return false;
  }

  rc = ssh_channel_open_session(shell_channel_);
  if (rc != SSH_OK)
    return false;

  rc = ssh_channel_request_shell(shell_channel_);
  if (rc != SSH_OK)
    return false;

  return true;
}

void SSH::ExecuteShell(const char *commandline) {
  std::string command = commandline;
  command.append("\n");
  ssh_channel_write(shell_channel_, command.c_str(), command.size());
}

bool SSH::ExecuteCommand(const char *commandline) {
  ssh_channel channel = ssh_channel_new(session_);
  if (channel == NULL) {
    return false;
  }

  int rc = ssh_channel_open_session(channel);
  if (rc < 0) {
    return false;
  }

  rc = ssh_channel_request_exec(channel, commandline);
  if (rc < 0) {
    return false;
  }

  char buffer[256];
  int nbytes = ssh_channel_read(channel, buffer, sizeof(buffer), 0);
  while (nbytes > 0) {
    nbytes = ssh_channel_read(channel, buffer, sizeof(buffer), 0);
  }

  if (nbytes < 0) {
    return false;
  }

  ssh_channel_send_eof(channel);
  ssh_channel_close(channel);
  ssh_channel_free(channel);
  return true;
}

bool SSH::ExecuteCommand(const char *commandline, std::string &result) {
  ssh_channel channel = ssh_channel_new(session_);
  if (channel == NULL) {
    return false;
  }

  int rc = ssh_channel_open_session(channel);
  if (rc < 0) {
    return false;
  }

  rc = ssh_channel_request_exec(channel, commandline);
  if (rc < 0) {
    return false;
  }

  char buffer[256];
  int nbytes = ssh_channel_read(channel, buffer, sizeof(buffer), 0);
  while (nbytes > 0) {
    result.append(buffer, nbytes);
    nbytes = ssh_channel_read(channel, buffer, sizeof(buffer), 0);
  }

  if (nbytes < 0) {
    return false;
  }

  ssh_channel_send_eof(channel);
  ssh_channel_close(channel);
  ssh_channel_free(channel);
  return true;
}

bool SSH::Download(const char *remote_path, const char *local_path) {
  char buffer[16384];
  ssh_scp scp = ssh_scp_new(session_,
                            SSH_SCP_READ | SSH_SCP_RECURSIVE, remote_path);
  if (ssh_scp_init(scp) != SSH_OK) {
    fprintf(stderr, "error initializing scp: %s\n", ssh_get_error(session_));
    ssh_scp_free(scp);
    return false;
  }
  std::string dirpath = local_path;
  do {
    int rc = ssh_scp_pull_request(scp);
    switch (rc) {
    case SSH_SCP_REQUEST_NEWFILE: {
      int size = ssh_scp_request_get_size(scp);
      const char *filename = ssh_scp_request_get_filename(scp);
      int mode = ssh_scp_request_get_permissions(scp);
      if (ssh_scp_accept_request(scp) == SSH_ERROR) {
        ssh_scp_close(scp);
        ssh_scp_free(scp);
        return false;
      }

      std::string filepath = Format("%s/%s", dirpath.c_str(), filename);
      FILE *file = fopen(filepath.c_str(), "wb");
      int total = 0;
      while (true) {
        rc = ssh_scp_read(scp, buffer, sizeof(buffer));
        if (rc == SSH_ERROR) {
          fprintf(stderr, "Error reading scp: %s\n", ssh_get_error(session_));
          ssh_scp_close(scp);
          ssh_scp_free(scp);
          return false;
        }

        fwrite(buffer, rc, 1, file);
        total += rc;
        if (total == size) {
          fclose(file);
          file = NULL;
          break;
        }
      }
      break;
    }
    case SSH_ERROR:
      fprintf(stderr, "Error: %s\n", ssh_get_error(session_));
      ssh_scp_close(scp);
      ssh_scp_free(scp);
      return false;
    case SSH_SCP_REQUEST_WARNING:
      fprintf(stderr, "Warning: %s\n", ssh_scp_request_get_warning(scp));
      break;
    case SSH_SCP_REQUEST_NEWDIR: {
      const char *dirname = ssh_scp_request_get_filename(scp);
      int mode = ssh_scp_request_get_permissions(scp);
      ssh_scp_accept_request(scp);
      dirpath += "/";
      dirpath += dirname;
      ssh_mkdir(dirpath.c_str(), 0755);
      break;
    }
    case SSH_SCP_REQUEST_ENDDIR: {
      size_t index = dirpath.find_last_of('/');
      dirpath = dirpath.substr(0, index);
      break;
    }
    case SSH_SCP_REQUEST_EOF:
      printf("End of requests\n");
      goto end;
    }
  } while (true);
end:
  ssh_scp_close(scp);
  ssh_scp_free(scp);
  return true;
}

bool SSH::Upload(const char* local_path, const char *remote_path) {
  ssh_scp scp = NULL;
  FILE *file = NULL;
  char *filename = NULL;

  scp = ssh_scp_new(session_, SSH_SCP_WRITE, remote_path);
  if (ssh_scp_init(scp) != SSH_OK) {
    fprintf(stderr, "error initializing scp: %s\n", ssh_get_error(session_));
    goto error;
  }

  file = fopen(local_path, "rb");
  if (file == NULL) {
    goto error;
  }

  int fd = _fileno(file);

  if (fd < 0) {
    goto error;
  }

  struct stat s;
  int rc = fstat(fd, &s);
  if (rc != 0) {
    goto error;
  }
  int size = s.st_size;
  int mode = s.st_mode & ~S_IFMT;
  filename = ssh_basename(local_path);
  rc = ssh_scp_push_file(scp, filename, size, mode);
  if (rc == SSH_ERROR) {
    goto error;
  }

  int total = 0;
  char buffer[16384];
  do {
    memset(buffer, 0, sizeof(buffer));
    size_t read = fread(buffer, 1, sizeof(buffer), file);
    if (read == 0) break;
    if (read < 0) {
      goto error;
    }

    int write = ssh_scp_write(scp, buffer, read);
    if (write == SSH_ERROR) {
      goto error;
    }
    total += read;
  } while (total < size);

  fclose(file);
  ssh_string_free_char(filename);
  rc = ssh_scp_close(scp);
  if (rc == SSH_ERROR) {
    ssh_scp_free(scp);
    return false;
  }
  return true;

error:
  if (file != NULL) {
    fclose(file);
    file = NULL;
  }
  if (filename != NULL) {
    ssh_string_free_char(filename);
    filename = NULL;
  }
  if (scp != NULL) {
    ssh_scp_free(scp);
    scp = NULL;
  }
  return false;
}

bool SSH::Upload(const char* data,
                 size_t data_size,
                 const char* filename,
                 const char* remote_path) {
  ssh_scp scp = ssh_scp_new(session_, SSH_SCP_WRITE, remote_path);
  if (ssh_scp_init(scp) != SSH_OK) {
    fprintf(stderr, "error initializing scp: %s\n", ssh_get_error(session_));
    ssh_scp_free(scp);
    return false;
  }

  int rc = ssh_scp_push_file(scp, filename, data_size, 0644);
  if (rc == SSH_ERROR) {
    ssh_scp_free(scp);
    return false;
  }

  int write = ssh_scp_write(scp, data, data_size);
  if (write == SSH_ERROR) {
    ssh_scp_free(scp);
    return false;
  }

  rc = ssh_scp_close(scp);
  if (rc == SSH_ERROR) {
    ssh_scp_free(scp);
    return false;
  }
  return true;
}

bool SSH::IsExist(const char *commandline, const char *find) {
  std::string result;
  if (!ExecuteCommand(commandline, result)) return false;
  return (std::string::npos != result.find(find));
}

void SSH::Uninitialize() {
  if (shell_channel_ != NULL) {
    ssh_channel_send_eof(shell_channel_);
    ssh_channel_close(shell_channel_);
    ssh_channel_free(shell_channel_);
  }


  ssh_disconnect(session_);
  ssh_free(session_);
  ssh_finalize();
}
